from .model import ModelBase, message_to_str
from .model import ModelBase, Message, StarCoder
from tracing import get_code_traces_block, get_code_traces_line, get_code_traces_function
from typing import Optional, List, Union, Any
import ast
import re, random, time, json
import logging
from .parse import parse_code_block, add_code_block
from .prompt import *
from utils import IMPORT_HEADER

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('prompt_logs.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def log_prompt(messages: Union[List[Message], str], prefix: str = ""):
    """记录传递给模型的 prompt"""
    logger.info(f"\n{'='*50}\n{prefix} PROMPT:\n{'='*50}")
    if isinstance(messages, list):
        for msg in messages:
            logger.info(f"Role: {msg.role}\nContent:\n{msg.content}\n{'-'*50}")
    else:
        logger.info(f"Content:\n{messages}\n{'-'*50}")
    logger.info(f"{'='*50}\n")

def trim_header(func_impl):
    if IMPORT_HEADER in func_impl:
        func_impl = func_impl.replace(IMPORT_HEADER, "")
    return func_impl

def parse_explanation(responses, trace_blocks, prev_func_impl):
    lines = responses.strip().strip('.').split("\n")
    wrong_block = None
    explanation = None
    for l in lines:
        try:
            item = json.loads(l)
            assert "block" in item
            assert "correct" in item
            assert "explanation" in item
        except Exception as e:
            continue
        # convert parsed_reponse['correct'] into bool
        if isinstance(item['correct'], bool):
            item['correct'] = item['correct']
        elif isinstance(item['correct'], str):
            item['correct'] = item['correct'].lower() == 'true'
        else:
            assert False, "Strange type of correct field"
        # Check correct
        if not item['correct']:
            block_id = int(item["block"].replace("BLOCK-", ""))
            if block_id < len(trace_blocks):
                wrong_block = "\n".join(trace_blocks[block_id])
            explanation = item['explanation']
    if wrong_block is None:
        wrong_block = "\n".join([(f"[BLOCK-{i}]\n" + "\n".join(b)) for i, b in enumerate(trace_blocks)])
    if explanation is None:
        explanation = responses
    return wrong_block, explanation

def print_messages(messages: List[Message], prefix = "") -> None:
    print("::CHAT MESSAGE::" +prefix)
    for msg in messages:
        print(msg.content)
    print("==================")

def parse_debug_response(text):
    text = text.replace('```json', '').replace('```', '')
    text = text.replace('\\', '\\\\').replace("\'correct\'", "\"correct\"").replace("\'explanation\'", "\"explanation\"").replace(" \'", " \"").replace("\'}", "\"}").strip()
    assert text[0] == '{'
    if text[-1] != '}':
        if text[-1] in ["\"", "\'"]:
            text += '}'
        else:
            text += "\"}"
    text = text.replace("\'", "\"")
    text = text.replace("\"", "\\\"")
    text = text.replace("\\\"correct\\\"", "\"correct\"")
    text = text.replace("\\\"explanation\\\"", "\"explanation\"")
    text = text.replace(": \\\"", ": \"")
    text = text.replace("\\\"", "\"")
    text = text.replace("\\\"}", "\"}")
    text = text.replace('False', 'false').replace('True', 'true')
    text = text.replace(": false,", ": \"false\",")
    try:
        json_item = json.loads(text, strict=False)
    except Exception as e:
        return {"correct": False, "explanation": "I don't know why this is wrong but it is incorrect."}
    return json_item

def check_block_correctness(messages, model: ModelBase,failed_test: str, block: str):
    if model.is_chat:
        debug_message = [
            Message(
                role="user",
                content=f"### Debug Start\n## Execution Trace\n{block}\n### Debug Response"
            )
        ]
        messages += debug_message
        max_trials = 10
        trials = 0
        debug_response = None
        while trials < max_trials:
            try:
                print_messages(debug_message)
                debug_response = model.generate_chat(messages=messages, stop=["### Debug End"])
            except Exception as e:
                print("Generate Debug Response Fail:\n" + str(e))
                trials += 1
                time.sleep(5)
                continue
            else:
                break
        if debug_response is None:
            debug_response = "{\"correct\": \"false\", \"explanation\": \"I don't know why this is wrong.\"}"
        debug_response = debug_response.strip()
        print(debug_response+"\n### Debug End")
    else:
        messages += f"### Debug Start\n### Execution Trace\n{block}\n### Debug Response"
        debug_response = model.generate_completion(messages, temperature=0, stop=["### Debug End"])
    parsed_response = parse_debug_response(debug_response)
    # Update messages
    if model.is_chat:
        messages.append(Message(role="assistant", content=debug_response + "\n### Debug End"))
    else:
        messages += debug_response + "### Debug End\n"
    # convert parsed_reponse['correct'] into bool
    if isinstance(parsed_response['correct'], bool):
        is_correct = parsed_response['correct']
    elif isinstance(parsed_response['correct'], str):
        is_correct = parsed_response['correct'].lower() == 'true'
    else:
        assert False, "Strange type of correct field"
    if 'explanation' not in parsed_response:
        parsed_response['explanation'] = "I don't know why this is wrong."
    return is_correct, parsed_response['explanation'], messages

def get_code_body(response):
    if "```" in response:
        return response.split("```python")[1].split("```")[0]
    else:
        return response

class PyGenerator:
    def ldb_debug(
        self,
        prompt: str,
        prev_func_impl: str,
        failed_test: str,
        entry: Optional[str],
        model: ModelBase,
        messages: Union[List[Message], str],
        dataset_type: str = "",
        level: str = "block",
    ) -> Union[List[Message], str]:
        """调试函数实现"""
        # 记录初始参数
        logger.info(f"\n{'='*50}\nDEBUG PARAMETERS:\n{'='*50}")
        logger.info(f"Prompt: {prompt}")
        logger.info(f"Failed Test: {failed_test}")
        logger.info(f"Entry: {entry}")
        logger.info(f"Dataset Type: {dataset_type}")
        logger.info(f"Level: {level}")
        logger.info(f"{'='*50}\n")
        
        # 获取失败的测试用例
        failed_test_string = failed_test
        real_test_output = ""
        
        # 根据数据集类型处理测试用例
        if dataset_type == "ACM":
            # ACM 模式下的测试数据格式化
            if len(messages) == 0:
                messages = [
                    Message(
                        role = "system",
                        content = "You are an expert programming assistant.",
                    ),
                    Message(
                        role = "user",
                        content = f"Complete the following task in Python. Please respond with code only (with the code inside a Markdown code block).\n{prompt}"
                    ),
                    Message(
                        role = "assistant",
                        content = f"{prev_func_impl}"
                    )
                ]
                print_messages(messages, "ACM:\n")
                log_prompt(messages, "ACM INITIAL")
            
            # 确保测试数据是纯文本格式
            if isinstance(failed_test, dict):
                input_data = failed_test.get('input', '')
                output_data = failed_test.get('output', '')
                actual_output = failed_test.get('actual_output', None)
                error = failed_test.get('error', '')
                
                # 确保输入和输出都是字符串
                input_str = str(input_data).strip()
                output_str = str(output_data).strip()
                actual_output_str = str(actual_output).strip() if actual_output else 'null'
                error_str = str(error).strip() if error else ''
                
                failed_test_string = f"Input:\n{input_str}\nExpected Output:\n{output_str}"
                if actual_output_str != 'null':
                    failed_test_string += f"\nActual Output:\n{actual_output_str}"
                if error_str:
                    failed_test_string += f"\nError:\n{error_str}"
                
                logger.info(f"\n{'='*50}\nFORMATTED TEST DATA:\n{'='*50}")
                logger.info(failed_test_string)
                logger.info(f"{'='*50}\n")
            
            # 构建调试 prompt
            feedback = f"The code above fails the given test case:\n{failed_test_string}. \nHelp me debug this.\n"
            log_prompt([Message(role="user", content=feedback)], "ACM FEEDBACK")
            
            # 获取追踪信息
            logger.info(f"\n{'='*50}\nGETTING TRACE BLOCKS:\n{'='*50}")
            logger.info(f"Level: {level}")
            logger.info(f"Program Code:\n{IMPORT_HEADER + prev_func_impl}")
            logger.info(f"Test Code:\n{failed_test_string}")
            
            if level == "line":
                trace_blocks = get_code_traces_line(IMPORT_HEADER + prev_func_impl, failed_test_string, is_acm=True)
            elif level == "function":
                trace_blocks = get_code_traces_function(IMPORT_HEADER + prev_func_impl, failed_test_string, is_acm=True)
            else:
                trace_blocks = get_code_traces_block(IMPORT_HEADER + prev_func_impl, failed_test_string, None, is_acm=True)
            
            logger.info(f"\n{'='*50}\nTRACE BLOCKS RESULT:\n{'='*50}")
            logger.info(f"Trace Blocks Type: {type(trace_blocks)}")
            logger.info(f"Trace Blocks Content:\n{trace_blocks}")
            
            # 处理追踪块
            if isinstance(trace_blocks, str):
                if trace_blocks == "*timeout*":
                    real_test_output = "Timeout"
                elif trace_blocks.startswith("*execution fail*"):
                    real_test_output = trace_blocks
                elif trace_blocks.startswith("*parse fail*"):
                    real_test_output = trace_blocks
            else:
                # 格式化 trace blocks
                trace_str = ""
                for i, block in enumerate(trace_blocks):
                    trace_str += f"Block {i+1}:\n"
                    for line in block:
                        # 确保 line 是字符串类型
                        if not isinstance(line, str):
                            line = str(line)
                        trace_str += line + "\n"
                    trace_str += "\n"
                
                # 添加错误信息
                if failed_test:
                    trace_str += f"Failed Test:\n{failed_test}\n\n"
                
                # 添加调试消息
                trace_str += "Debug Message:\n"
                trace_str += "Please analyze the execution trace and identify the issue in the code. "
                trace_str += "Focus on the variable values and control flow to find the bug.\n"
                
                real_test_output = trace_str
            
            logger.info(f"\n{'='*50}\nFORMATTED TRACE OUTPUT:\n{'='*50}")
            logger.info(real_test_output)
            
            # 添加追踪结果到消息列表
            messages.append(Message(role="user", content=feedback))
            messages.append(Message(role="user", content=f"Here is the execution trace:\n{real_test_output}\nPlease help me fix the code."))
            
            logger.info(f"\n{'='*50}\nFINAL MESSAGES:\n{'='*50}")
            for msg in messages:
                logger.info(f"Role: {msg.role}")
                logger.info(f"Content:\n{msg.content}\n{'-'*50}")
            
            # 获取模型的响应
            response = model.generate_chat(messages=messages, num_comps=1, temperature=0)
            messages.append(Message(role="assistant", content=response))
            
            return messages
        elif dataset_type in ["TransCoder"]:
            if len(messages) == 0:
                messages = [
                    Message(
                        role = "system",
                        content = "You are an expert programming assistant.",
                    ),
                    Message(
                        role = "user",
                        content = f"Translate the following C++ code into Python. Please respond with code only (with the code inside a Markdown code block).\n{prompt}"
                    ),
                    Message(
                        role = "assistant",
                        content = f"{prev_func_impl}"
                    )
                ]
                print_messages(messages, "213:\n")
            feedback = f"The Python translation does not do the same thing as the C++ code. Help me debug this. \nThese are the results of one failed unit test that tests whether the Python translation's outputs match the C++ program's outputs:\n{failed_test}."
        elif dataset_type in ["HumanEval", "MBPP"]:
            if len(messages) == 0:
                messages = [
                    Message(
                        role = "system",
                        content = "You are an expert programming assistant.",
                    ),
                    Message(
                        role = "user",
                        content = f"Complete the following task in Python. Please respond with code only (with the code inside a Markdown code block).\n{prompt}"
                    ),
                    Message(
                        role = "assistant",
                        content = f"{prev_func_impl}"
                    )
                ]
                print_messages(messages, "268:\n")
            feedback = f"The code above fails the given unit test:\n{failed_test}. \nHelp me debug this.\n"
        
        # 添加调试请求到消息列表
        messages.append(Message(role="user", content=feedback))
        
        # 获取模型的响应
        response = model.generate_chat(messages=messages, num_comps=1, temperature=0)
        messages.append(Message(role="assistant", content=response))
        
        # 从响应中提取代码
        code = get_code_body(response)
        if not code:
            return messages
        
        # 获取追踪信息
        if level == "line":
            trace_blocks = get_code_traces_line(IMPORT_HEADER + prev_func_impl, failed_test_string, None, is_acm=True)
        elif level == "function":
            trace_blocks = get_code_traces_function(IMPORT_HEADER + prev_func_impl, failed_test_string, None, is_acm=True)
        else:
            trace_blocks = get_code_traces_block(IMPORT_HEADER + prev_func_impl, failed_test_string, None, is_acm=True)
        
        print("Get trace blocks...")
        
        # 处理追踪块
        if isinstance(trace_blocks, str):
            if trace_blocks == "*timeout*":
                real_test_output = "Timeout"
            elif trace_blocks.startswith("*execution fail*"):
                real_test_output = trace_blocks
            elif trace_blocks.startswith("*parse fail*"):
                real_test_output = trace_blocks
        else:
            # 格式化 trace blocks
            trace_str = ""
            for i, block in enumerate(trace_blocks):
                trace_str += f"Block {i+1}:\n"
                for line in block:
                    # 确保 line 是字符串类型
                    if not isinstance(line, str):
                        line = str(line)
                    trace_str += line + "\n"
                trace_str += "\n"
            
            # 添加错误信息
            if failed_test:
                trace_str += f"Failed Test:\n{failed_test}\n\n"
            
            # 添加调试消息
            trace_str += "Debug Message:\n"
            trace_str += "Please analyze the execution trace and identify the issue in the code. "
            trace_str += "Focus on the variable values and control flow to find the bug.\n"
            
            real_test_output = trace_str
        
        # 添加追踪结果到消息列表
        messages.append(Message(role="user", content=f"Here is the execution trace:\n{real_test_output}\nPlease help me fix the code."))
        response = model.generate_chat(messages=messages, num_comps=1, temperature=0)
        messages.append(Message(role="assistant", content=response))
        
        return messages

    def ldb_generate(
        self,
        prompt: str,
        model: ModelBase,
        messages: Union[List[Message], str],
        dataset_type: str = "",
    ) -> Union[List[Message], str]:
        """生成函数实现"""
        # 记录初始参数
        logger.info(f"\n{'='*50}\nGENERATE PARAMETERS:\n{'='*50}")
        logger.info(f"Prompt: {prompt}")
        logger.info(f"Dataset Type: {dataset_type}")
        logger.info(f"{'='*50}\n")
        
        # 根据数据集类型处理生成请求
        if dataset_type == "ACM":
            # ACM 模式下的生成请求
            if len(messages) == 0:
                messages = [
                    Message(
                        role = "system",
                        content = "You are an expert programming assistant.",
                    ),
                    Message(
                        role = "user",
                        content = f"Complete the following task in Python. Please respond with code only (with the code inside a Markdown code block).\n{prompt}"
                    )
                ]
                print_messages(messages, "ACM:\n")
                log_prompt(messages, "ACM GENERATE")
        elif dataset_type in ["TransCoder"]:
            if len(messages) == 0:
                messages = [
                    Message(
                        role = "system",
                        content = "You are an expert programming assistant.",
                    ),
                    Message(
                        role = "user",
                        content = f"Translate the following C++ code into Python. Please respond with code only (with the code inside a Markdown code block).\n{prompt}"
                    )
                ]
                print_messages(messages, "213:\n")
        elif dataset_type in ["HumanEval", "MBPP"]:
            if len(messages) == 0:
                messages = [
                    Message(
                        role = "system",
                        content = "You are an expert programming assistant.",
                    ),
                    Message(
                        role = "user",
                        content = f"Complete the following task in Python. Please respond with code only (with the code inside a Markdown code block).\n{prompt}"
                    )
                ]
                print_messages(messages, "268:\n")
        
        # 生成代码
        if model.is_chat:
            response = model.generate_chat(messages=messages, num_comps=1, temperature=0)
            messages.append(Message(role="assistant", content=response))
            print_messages([Message(role="assistant", content=response)])
            return messages
        else:
            if dataset_type in ["TransCoder"]:
                if len(messages) == 0:
                    messages = f"{PY_CHAINOFDEBUG_TRANSLATION_INSTRUCTION}"
                    print(messages)
                    # Explain C++
                    delta_msg = f"\n[c++]\n{self.get_last_cpp(prompt)}\n[/c++]\n[explanation]"
                    print(delta_msg)
                    messages += delta_msg
                    explanation = model.generate_completion(messages, temperature=0, stop=["[/explanation]"])
                    delta_msg = f"\n{explanation.strip()}\n[/explanation]\n[python]"
                    print(delta_msg)
                    messages += delta_msg
                    response = model.generate_completion(messages, temperature=0, stop=["[/python]"])
                    delta_msg = f"\n{response.strip()}\n[/python]"
                    print(delta_msg)
                    messages += delta_msg
                    
                    # 将字符串消息转换为消息列表
                    return [
                        Message(role="system", content="You are an expert programming assistant."),
                        Message(role="user", content=messages),
                        Message(role="assistant", content=response)
                    ]
            else:
                if len(messages) == 0:
                    messages = f"{PY_CHAINOFDEBUG_TEXT2CODE_INSTRUCTION}\n{prompt}\n"
                    print(messages)
                else:
                    # 将消息列表转换为字符串
                    if isinstance(messages, list):
                        messages_str = ""
                        for msg in messages:
                            if msg.role == "user":
                                messages_str += f"User: {msg.content}\n"
                            elif msg.role == "assistant":
                                messages_str += f"Assistant: {msg.content}\n"
                            elif msg.role == "system":
                                messages_str += f"System: {msg.content}\n"
                        messages = messages_str
                    
                    delta_msg = f"### Task Start ###\n{prompt}\n"
                    messages += delta_msg
                    print(delta_msg)
                response = model.generate_completion(messages, temperature=0, stop=["### Task End ###"])
                delta_msg = f"\n{response.strip()}\n### Task End ###"
                print(delta_msg)
                messages += delta_msg
                
                # 将字符串消息转换为消息列表
                return [
                    Message(role="system", content="You are an expert programming assistant."),
                    Message(role="user", content=messages),
                    Message(role="assistant", content=response)
                ]

    def get_last_cpp(self, prompt):
        return prompt.split("[c++]\n")[-1].replace("\n[python]", "")

    def simple_translation(self, func_sig, model, prev_func_impl, feedback, given_tests, num_comps, temperature):
        assertion_string = "\n".join(given_tests)
        if len(assertion_string) > 3000:
            assertion_string = "\n".join(given_tests[:5])
        if model.is_chat:
            system_prompt = "You are an expert programming assistant."
            user_prompt = f"Translate the C++ code into Python code. Please respond with code only (with the code inside a Markdown code block). These are the assertions for your function for your reference. Answer with code only:\n{assertion_string}\n{func_sig}"
            print(system_prompt + "\n" + user_prompt)
            messages = [
                Message(
                    role="system",
                    content=system_prompt,
                ),
                Message(
                    role="user",
                    content=user_prompt,
                ),
            ]
            func_bodies = model.generate_chat(messages=messages, num_comps=num_comps, temperature=0)
        else:
            messages = f"Translate the following C++ program into Python\n{func_sig}"
            func_bodies = model.generate_completion(messages, temperature=0, stop=["[c++]", "[/code]"])
        return func_bodies

    def simple_text2code(self, func_sig, model, prev_func_impl, feedback, given_tests, num_comps, temperature):
        if model.is_chat:
            func_sig = func_sig.rstrip('\n')
            user_prompt = f"Complete the following task in Python. Remember to repeat all imports and function header. Here is a unit test:\n{given_tests[0].strip()}\n\n{func_sig}"
            messages = [
                Message(
                    role="system",
                    content=f"You are an expert programming assistant.",
                ),
                Message(
                    role="user",
                    content=user_prompt,
                ),
            ]
            func_bodies = model.generate_chat(messages=messages, num_comps=num_comps, temperature=0)
        else:
            messages = f"# Write Python function to complete the task and pass the assertion tests.\n\n### Task Start ###\n# These are the assertions for your function:\nassert similar_elements((3, 4, 5, 6),(5, 7, 4, 10)) == (4, 5)\n\ndef similar_elements(test_tup1, test_tup2):\n\"\"\" Write a function to find the similar elements from the given two tuple lists. \"\"\"\n    res = tuple(set(test_tup1) & set(test_tup2))\n    return (res)\n### Task End ###\n\n### Task Start ###\n# These are the assertions for your function:\nassert is_not_prime(2) == False\n\nimport math\ndef is_not_prime(n):\n    \"\"\" Write a python function to identify non-prime numbers. \"\"\"\n    result = False\n    for i in range(2,int(math.sqrt(n)) + 1):\n        if n % i == 0:\n            result = True\n    return result\n### Task End ###\n\n### Task Start ###\n# These are the assertions for your function:\nassert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],3)==[85, 75, 65]\n\nimport heapq as hq\ndef heap_queue_largest(nums,n):\n    \"\"\" Write a function to find the largest integers from a given list of numbers using heap queue algorithm. \"\"\"\n    largest_nums = hq.nlargest(n, nums)\n    return largest_nums\n### Task End ###\n\n### Task Start ###\n# These are the assertions for your function:\n{given_tests[0].strip()}\n\n{func_sig.strip()}"
            print(messages)
            func_bodies = model.generate_completion(messages, temperature=0, stop=["### Task End ###"])
        return func_bodies, messages

    def func_impl(
        self,
        func_sig: str,
        model: ModelBase,
        strategy: str,
        prev_func_impl: Optional[str] = None,
        feedback: Optional[str] = None,
        given_tests: Optional[str] = None,
        num_comps: int = 1,
        temperature: float = 0.0,
        dataset_type: str = "",
        prompt: str = ""
    ) -> Union[str, List[str]]:
        # Validate
        if strategy not in ["simple"]:
            raise ValueError(
                f"Invalid strategy: given `{strategy}` but expected `simple`")
        if model.is_chat:
            if strategy == "simple":
                # Translation Task
                if dataset_type in ["TransCoder"]:
                    func_bodies, messages = self.simple_translation(func_sig, model, prev_func_impl, feedback, given_tests, num_comps, temperature)
                else:
                    func_bodies, messages = self.simple_text2code(func_sig, model, prev_func_impl, feedback, given_tests, num_comps, temperature)
            else:
                assert False, "Not Impl!"
        else:
            if strategy == "simple":
                # Translation Task
                messages = "" # placeholder
                if dataset_type in ["TransCoder"]:
                    func_bodies = self.simple_translation(func_sig, model, prev_func_impl, feedback, given_tests, num_comps, temperature)
                else:
                    func_bodies, messages = self.simple_text2code(func_sig, model, prev_func_impl, feedback, given_tests, num_comps, temperature)
            else:
                assert False, "Not Impl!"
        
        if num_comps == 1:
            assert isinstance(func_bodies, str)
            func_body_str = get_code_body(func_bodies)
            if isinstance(messages, list):
                if strategy == 'simple':
                    messages.append(Message(role="assistant", content=func_bodies))
                else:    
                    messages.append(Message(role="assistant", content=func_body_str))
            elif isinstance(messages, str):
                messages += "\n" + func_body_str
            else:
                assert False, "Not Impl!"
            return func_body_str, messages
        else:
            messages += [Message(role="assistant", content=func_body) for func_body in func_bodies]
            func_bodies = [get_code_body(func_body) for func_body in func_bodies]
            return func_bodies, messages

    def generate_chat_messages(self, dataset_type: str, prompt: str, messages: Union[List, str], entry_point: str, model: Any, level: str) -> List:
        """Generate chat messages based on dataset type."""
        if dataset_type == "TransCoder":
            user_message = f"{PY_CHAINOFDEBUG_TRANSLATION_INSTRUCTION}\n{prompt}"
        elif dataset_type == "HumanEval":
            user_message = f"{PY_CHAINOFDEBUG_TEXT2CODE_INSTRUCTION}\n{prompt}"
        elif dataset_type == "ACM":
            user_message = f"{PY_CHAINOFDEBUG_ACM_INSTRUCTION}\n{prompt}"
        else:
            raise ValueError(f"Invalid dataset type: {dataset_type}")

        if isinstance(messages, str):
            messages = [{"role": "user", "content": user_message}]
        else:
            messages.append({"role": "user", "content": user_message})

        return messages

DUMMY_FUNC_SIG = "def func():"
DUMMY_FUNC_CALL = "func()"

def handle_first_line_indent(func_body: str) -> str:
    if func_body.startswith("    "):
        return func_body
    split = func_body.splitlines()
    return f"    {split[0]}\n" + "\n".join(split[1:])

def handle_entire_body_indent(func_body: str) -> str:
    split = func_body.splitlines()
    res = "\n".join(["    " + line for line in split])
    return res

def fix_turbo_response(func_body: str) -> str:
    return fix_markdown(remove_unindented_signatures(func_body))

def fix_markdown(func_body: str) -> str:
    return re.sub("`{3}", "", func_body)

def remove_unindented_signatures(code: str) -> str:
    regex = r"^def\s+\w+\s*\("

    before_signature = []
    after_signature = []
    signature_found = False

    for line in code.split("\n"):
        if re.match(regex, line):
            signature_found = True
            continue

        if signature_found:
            after_signature.append(line)
        else:
            if not line.startswith("    ") and line.strip():
                line = "    " + line
            before_signature.append(line)

    return "\n".join(before_signature + after_signature)

def py_fix_indentation(func_body: str) -> str:
    func_body = fix_turbo_response(func_body)
    """
    3 cases:
        1. good syntax
        2. first line not good
        3. entire body not good
    """
    def parse_indent_rec(f_body: str, cur_state: int) -> str:
        f_body = fix_markdown(f_body)
        if cur_state > 1:
            return f_body
        code = f'{DUMMY_FUNC_SIG}\n{f_body}\n{DUMMY_FUNC_CALL}'
        try:
            exec(code)
            return f_body
        except (IndentationError, SyntaxError):
            p_func = handle_first_line_indent if cur_state == 0 else handle_entire_body_indent
            return parse_indent_rec(p_func(func_body), cur_state + 1)
        except Exception:
            return f_body
    return parse_indent_rec(func_body, 0)

def py_is_syntax_valid(code: str) -> bool:
    try:
        ast.parse(code)
        return True
    except Exception:
        return False
